package pl.th.permission.plugin;

import static org.objectweb.asm.Opcodes.AASTORE;
import static org.objectweb.asm.Opcodes.ACC_PRIVATE;
import static org.objectweb.asm.Opcodes.ACC_PUBLIC;
import static org.objectweb.asm.Opcodes.ALOAD;
import static org.objectweb.asm.Opcodes.ANEWARRAY;
import static org.objectweb.asm.Opcodes.DUP;
import static org.objectweb.asm.Opcodes.GOTO;
import static org.objectweb.asm.Opcodes.IFEQ;
import static org.objectweb.asm.Opcodes.IF_ICMPNE;
import static org.objectweb.asm.Opcodes.ILOAD;
import static org.objectweb.asm.Opcodes.INVOKESPECIAL;
import static org.objectweb.asm.Opcodes.INVOKESTATIC;
import static org.objectweb.asm.Opcodes.INVOKEVIRTUAL;
import static org.objectweb.asm.Opcodes.RETURN;
import static pl.th.permission.plugin.Utils.STR_ARR_EQUAL_METHOD_DESC;
import static pl.th.permission.plugin.Utils.STR_ARR_EQUAL_METHOD_NAME;
import static pl.th.permission.plugin.Utils.STR_ARR_EQUAL_METHOD_OWNER;

import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.ClassNode;
import org.objectweb.asm.tree.MethodNode;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import pl.th.permission.plugin.node.PermissionMethodNode;
import pl.th.permission.plugin.node.annotation.AllAnnotation;
import pl.th.permission.plugin.node.annotation.AnnType;
import pl.th.permission.plugin.node.annotation.PartAnnotation;

public class ClassInjector {

    private final static String METHOD_PARAMETER_1 = "requestCode";
    private final static String METHOD_PARAMETER_1_DESC = "I";
    private final static String METHOD_PARAMETER_2 = "permissions";
    private final static String METHOD_PARAMETER_2_DESC = "[Ljava/lang/String;";

    private String clzName;
    private final Map<Integer, PermissionMethodNode> needMap;
    private final Map<Integer, PermissionMethodNode> allGrantedMap;
    private final Map<Integer, PermissionMethodNode> allDeniedMap;
    private final Map<Integer, PermissionMethodNode> anyGrantedMap;
    private final Map<Integer, PermissionMethodNode> anyDeniedMap;
    private final Map<Integer, List<PermissionMethodNode>> onGrantedMap;
    private final Map<Integer, List<PermissionMethodNode>> onDeniedMap;

    private PermissionMethodNode grantedCallbackMethod;
    private PermissionMethodNode deniedCallbackMethod;

    {
        grantedCallbackMethod = null;
        deniedCallbackMethod = null;
        needMap = new LinkedHashMap<>(4);
        allGrantedMap = new LinkedHashMap<>(4);
        allDeniedMap = new LinkedHashMap<>(4);
        anyGrantedMap = new LinkedHashMap<>(4);
        anyDeniedMap = new LinkedHashMap<>(4);
        onGrantedMap = new LinkedHashMap<>(4);
        onDeniedMap = new LinkedHashMap<>(4);
    }

    public ClassInjector(String clzName, final List<MethodNode> methodNodes) {
        this.clzName = clzName.replace(File.separator, ".");
        for (MethodNode methodNode : methodNodes) {
            final PermissionMethodNode pmNode = (PermissionMethodNode) methodNode;
            if (pmNode.getNeedAnnotation() != null) {
                final int requestCode = pmNode.getNeedAnnotation().getRequestCode();
                if (needMap.getOrDefault(requestCode, null) != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.NEED_PERMISSION, requestCode);
                }
                needMap.put(requestCode, pmNode);
            }
            if (pmNode.getAllGrantedAnnotation() != null) {
                final int requestCode = pmNode.getAllGrantedAnnotation().getRequestCode();
                if (allGrantedMap.getOrDefault(requestCode, null) != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_ALL_PERMISSION_GRANTED, requestCode);
                }
                allGrantedMap.put(requestCode, pmNode);
            }
            if (pmNode.getAllDeniedAnnotation() != null) {
                final int requestCode = pmNode.getAllDeniedAnnotation().getRequestCode();
                if (allDeniedMap.getOrDefault(requestCode, null) != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_ALL_PERMISSION_DENIED, requestCode);
                }
                allDeniedMap.put(requestCode, pmNode);
            }
            if (pmNode.getAnyGrantedAnnotation() != null) {
                final int requestCode = pmNode.getAnyGrantedAnnotation().getRequestCode();
                if (anyGrantedMap.getOrDefault(requestCode, null) != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_ANY_PERMISSION_GRANTED, requestCode);
                }
                anyGrantedMap.put(requestCode, pmNode);
            }
            if (pmNode.getAnyDeniedAnnotation() != null) {
                final int requestCode = pmNode.getAnyDeniedAnnotation().getRequestCode();
                if (anyDeniedMap.getOrDefault(requestCode, null) != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_ANY_PERMISSION_DENIED, requestCode);
                }
                anyDeniedMap.put(requestCode, pmNode);
            }
            if (pmNode.getOnGrantedAnnotation() != null) {
                final int requestCode = pmNode.getOnGrantedAnnotation().getRequestCode();
                onGrantedMap.putIfAbsent(requestCode, new ArrayList<>(2));
                onGrantedMap.get(requestCode).add(pmNode);
            }
            if (pmNode.getOnDeniedAnnotation() != null) {
                final int requestCode = pmNode.getOnDeniedAnnotation().getRequestCode();
                onDeniedMap.putIfAbsent(requestCode, new ArrayList<>(2));
                onDeniedMap.get(requestCode).add(pmNode);
            }
            if (pmNode.getGrantedCbAnnotation() != null) {
                if (grantedCallbackMethod != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_GRANTED_CALLBACK, 0);
                }
                grantedCallbackMethod = pmNode;
            }
            if (pmNode.getDeniedCbAnnotation() != null) {
                if (deniedCallbackMethod != null) {
                    throw MultiPleAnnotationException.obtain(AnnType.ON_DENIED_CALLBACK, 0);
                }
                deniedCallbackMethod = pmNode;
            }
        }
        checkGrantedCallback();
        checkDeniedCallback();
    }

    private void checkGrantedCallback() {
        if (grantedCallbackMethod != null && (!allGrantedMap.isEmpty()
                || !anyGrantedMap.isEmpty() || !onGrantedMap.isEmpty())) {
            String message = AnnType.ON_GRANTED_CALLBACK.toString();
            message += " can not be used with ";
            if (!allGrantedMap.isEmpty()) {
                message += ", " + AnnType.ON_ALL_PERMISSION_GRANTED;
            }
            if (!anyGrantedMap.isEmpty()) {
                message += ", " + AnnType.ON_ANY_PERMISSION_GRANTED;
            }
            if (!onGrantedMap.isEmpty()) {
                message += ", " + AnnType.ON_PERMISSION_GRANTED;
            }
            message += " in the class";
            throw new ConflictedAnnotationException(message);
        }

        for (Integer key : anyGrantedMap.keySet()) {
            if (allGrantedMap.getOrDefault(key, null) != null) {
                String message = AnnType.ON_ALL_PERMISSION_GRANTED + " and "
                        + AnnType.ON_ANY_PERMISSION_GRANTED
                        + " can not own the same requestCode : " + key
                        + "in the class: " + clzName;
                throw new ConflictedAnnotationException(message);
            }
            if (onGrantedMap.getOrDefault(key, null) != null) {
                String message = AnnType.ON_PERMISSION_GRANTED + " and "
                        + AnnType.ON_ANY_PERMISSION_GRANTED
                        + " can not own the same requestCode : " + key
                        + "in the class: " + clzName;
                throw new ConflictedAnnotationException(message);
            }
        }

        for (Map.Entry<Integer, List<PermissionMethodNode>> entry : onGrantedMap.entrySet()) {
            final int requestCode = entry.getKey();
            final PermissionMethodNode node = needMap.get(requestCode);
            Set<String> allPermissions = null;
            if (allGrantedMap.get(requestCode) != null && node != null) {
                allPermissions = node.getNeedAnnotation().getPermissions();
            }
            final List<PermissionMethodNode> nodeList = entry.getValue();
            final int size = nodeList.size();
            for (int i = 0; i < size; i++) {
                PermissionMethodNode node1 = nodeList.get(i);
                final Set<String> s1 = node1.getOnGrantedAnnotation().getPermissions();
                for (int j = i + 1; j < size; j++) {
                    PermissionMethodNode node2 = nodeList.get(j);
                    final Set<String> s2 = node2.getOnGrantedAnnotation().getPermissions();
                    if (Utils.isEqualSet(s1, s2)) {
                        String message = AnnType.ON_PERMISSION_GRANTED.toString();
                        message += " with permissions " + Utils.toString(s1)
                                + " can`t be used twice and more in the class: " + clzName;
                        throw new ConflictedAnnotationException(message);
                    }
                }
                if (allPermissions != null) {
                    if (Utils.isEqualSet(allPermissions, s1)) {
                        String message = AnnType.ON_PERMISSION_GRANTED.toString();
                        message += "with permissions " + Utils.toString(s1)
                                + " is Conflicted with " + AnnType.ON_ALL_PERMISSION_GRANTED
                                + " in the class: " + clzName;
                        throw new ConflictedAnnotationException(message);
                    }
                }
            }

        }

    }

    private void checkDeniedCallback() {
        if (deniedCallbackMethod != null && (!allDeniedMap.isEmpty()
                || !anyDeniedMap.isEmpty() || !onDeniedMap.isEmpty())) {
            String message = AnnType.ON_DENIED_CALLBACK.toString();
            message += " can not be used with ";
            if (!allDeniedMap.isEmpty()) {
                message += ", " + AnnType.ON_ALL_PERMISSION_DENIED;
            }
            if (!anyDeniedMap.isEmpty()) {
                message += ", " + AnnType.ON_ANY_PERMISSION_DENIED;
            }
            if (!onDeniedMap.isEmpty()) {
                message += ", " + AnnType.ON_PERMISSION_DENIED;
            }
            message += " in the class: " + clzName;
            throw new ConflictedAnnotationException(message);
        }

        for (Integer key : anyDeniedMap.keySet()) {
            if (allDeniedMap.getOrDefault(key, null) != null) {
                String message = AnnType.ON_ALL_PERMISSION_DENIED + " and "
                        + AnnType.ON_ANY_PERMISSION_DENIED
                        + " can not own the same requestCode : " + key
                        + " in the class: " + clzName;
                throw new ConflictedAnnotationException(message);
            }
            if (onDeniedMap.getOrDefault(key, null) != null) {
                String message = AnnType.ON_PERMISSION_DENIED + " and "
                        + AnnType.ON_ANY_PERMISSION_DENIED
                        + " can not own the same requestCode : " + key
                        + " in the class: " + clzName;
                throw new ConflictedAnnotationException(message);
            }
        }

        for (Map.Entry<Integer, List<PermissionMethodNode>> entry : onDeniedMap.entrySet()) {
            final int requestCode = entry.getKey();
            final PermissionMethodNode node = needMap.get(requestCode);
            Set<String> allPermissions = null;
            if (allDeniedMap.get(requestCode) != null && node != null) {
                allPermissions = node.getNeedAnnotation().getPermissions();
            }
            final List<PermissionMethodNode> nodeList = entry.getValue();
            final int size = nodeList.size();
            for (int i = 0; i < size; i++) {
                PermissionMethodNode node1 = nodeList.get(i);
                final Set<String> s1 = node1.getOnDeniedAnnotation().getPermissions();
                for (int j = i + 1; j < size; j++) {
                    PermissionMethodNode node2 = nodeList.get(j);
                    final Set<String> s2 = node2.getOnDeniedAnnotation().getPermissions();
                    if (Utils.isEqualSet(s1, s2)) {
                        String message = AnnType.ON_PERMISSION_DENIED.toString();
                        message += " with permissions " + Utils.toString(s1)
                                + " can`t be used twice and more in the class: " + clzName;
                        throw new ConflictedAnnotationException(message);
                    }
                }
                if (allPermissions != null) {
                    if (Utils.isEqualSet(allPermissions, s1)) {
                        String message = AnnType.ON_PERMISSION_DENIED.toString();
                        message += "with permissions " + Utils.toString(s1)
                                + " is Conflicted with " + AnnType.ON_ALL_PERMISSION_DENIED
                                + " in the class: " + clzName;
                        throw new ConflictedAnnotationException(message);
                    }
                }
            }

        }
    }

    public void injectGrantMethodToClass(ClassNode classNode) {
        final MethodNode methodNode = (MethodNode) classNode.visitMethod(ACC_PUBLIC,
                Utils.GRANT_METHOD_NAME, Utils.METHOD_DESC_CALLBACK, null, null);
        methodNode.visitParameter(METHOD_PARAMETER_1, 0);
        methodNode.visitParameter(METHOD_PARAMETER_2, 0);
        methodNode.visitCode();
        if (grantedCallbackMethod != null) {
            injectInterceptedMethod(methodNode, classNode, grantedCallbackMethod);
        } else {
            if (onGrantedMap.isEmpty() && anyGrantedMap.isEmpty() && allGrantedMap.isEmpty()) {
                injectEmptyBody(classNode, methodNode);
            } else {
                injectCallbackBody(classNode, methodNode, true);
            }
        }
    }

    public void injectDenyMethodToClass(ClassNode classNode) {
        final MethodNode methodNode = (MethodNode) classNode.visitMethod(ACC_PUBLIC,
                Utils.DENIED_METHOD_NAME, Utils.METHOD_DESC_CALLBACK, null, null);
        methodNode.visitParameter(METHOD_PARAMETER_1, 0);
        methodNode.visitParameter(METHOD_PARAMETER_2, 0);
        methodNode.visitCode();
        if (deniedCallbackMethod != null) {
            injectInterceptedMethod(methodNode, classNode, deniedCallbackMethod);
        } else {
            if (onDeniedMap.isEmpty() && anyDeniedMap.isEmpty() && allDeniedMap.isEmpty()) {
                injectEmptyBody(classNode, methodNode);
            } else {
                injectCallbackBody(classNode, methodNode, false);
            }
        }
    }

    private void injectInterceptedMethod(MethodVisitor methodVisitor, ClassNode classNode,
                                         PermissionMethodNode methodNode) {
        Label label0 = new Label();
        methodVisitor.visitLabel(label0);
        methodVisitor.visitVarInsn(ALOAD, 0);
        methodVisitor.visitVarInsn(ILOAD, 1);
        methodVisitor.visitVarInsn(ALOAD, 2);
        methodVisitor.visitMethodInsn(getInvokeOpCode(methodNode.access),
                classNode.name, methodNode.name, methodNode.desc, false);
        Label label1 = new Label();
        methodVisitor.visitLabel(label1);
        methodVisitor.visitInsn(RETURN);
        Label label2 = new Label();
        methodVisitor.visitLabel(label2);
        methodVisitor.visitLocalVariable("this", getClassDesc(classNode.name), null, label0, label1, 0);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_1, METHOD_PARAMETER_1_DESC, null, label0, label1, 1);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_2, METHOD_PARAMETER_2_DESC, null, label0, label1, 2);
        methodVisitor.visitMaxs(3, 3);
        methodVisitor.visitEnd();
    }

    //需要检查最大栈是否正确
    private void injectEmptyBody(ClassNode classNode, MethodVisitor methodVisitor) {
        Label label0 = new Label();
        methodVisitor.visitLabel(label0);
        methodVisitor.visitInsn(RETURN);
        Label label1 = new Label();
        methodVisitor.visitLabel(label1);
        methodVisitor.visitLocalVariable("this", getClassDesc(classNode.name), null, label0, label1, 0);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_1, METHOD_PARAMETER_1_DESC, null, label0, label1, 1);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_2, METHOD_PARAMETER_2_DESC, null, label0, label1, 2);
        methodVisitor.visitMaxs(0, 3);
        methodVisitor.visitEnd();
    }

    private void injectCallbackBody(ClassNode classNode, MethodVisitor methodVisitor, boolean isGranted) {
        final Map<Integer, PermissionMethodNode> needMethodMap = new HashMap<>(needMap);
        Map<Integer, PermissionMethodNode> anyMap = isGranted ? anyGrantedMap : anyDeniedMap;
        Map<Integer, PermissionMethodNode> allMap = isGranted ? allGrantedMap : allDeniedMap;
        Map<Integer, List<PermissionMethodNode>> onMap = isGranted ? onGrantedMap : onDeniedMap;
        for (Iterator<Map.Entry<Integer, PermissionMethodNode>> iterator = needMethodMap.entrySet().iterator(); iterator.hasNext(); ) {
            int tempCode = iterator.next().getKey();
            if (anyMap.get(tempCode) == null
                    && onMap.get(tempCode) == null
                    && allMap.get(tempCode) == null)
                iterator.remove();
        }
        int codeSize = needMethodMap.keySet().size();
        int index = 0;
        Label realStart = null;
        Label end = new Label();
        Label start;
        Label next = null;
        for (Integer requestCode : needMethodMap.keySet()) {
            if (index == 0) {
                start = new Label();
            } else {
                start = next;
            }
            if (index < codeSize - 1) {
                next = new Label();
            } else {
                next = end;
            }
            if (realStart == null) {
                realStart = start;
            }
            methodVisitor.visitLabel(start);
            if (index != 0) {
                methodVisitor.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
            }
            methodVisitor.visitVarInsn(ILOAD, 1);
            push(methodVisitor, requestCode);
            methodVisitor.visitJumpInsn(IF_ICMPNE, next);
            PermissionMethodNode anyMethodNode = anyMap.get(requestCode);
            if (anyMethodNode != null) {
                Label label0 = new Label();
                methodVisitor.visitLabel(label0);
                methodVisitor.visitVarInsn(ALOAD, 0);
                methodVisitor.visitVarInsn(ALOAD, 2);
                methodVisitor.visitMethodInsn(getInvokeOpCode(anyMethodNode.access), classNode.name, anyMethodNode.name, anyMethodNode.desc, false);
                if (index != codeSize - 1) {
                    methodVisitor.visitJumpInsn(GOTO, end);
                }
            } else {
                List<PermissionMethodNode> onMethodNodeList = onMap.get(requestCode);
                if (onMethodNodeList == null) onMethodNodeList = new LinkedList<>();
                PermissionMethodNode allMapNode = allMap.get(requestCode);
                if (allMapNode != null) {
                    onMethodNodeList.add(0, allMapNode);
                }
                int infoSize = onMethodNodeList.size();
                Label subStart;
                Label subNext = null;
                for (int i = 0; i < infoSize; i++) {
                    PermissionMethodNode tempNode = onMethodNodeList.get(i);
                    Set<String> permissions;
                    if (i == 0) {
                        final AllAnnotation allAnnotation = isGranted
                                ? tempNode.getAllGrantedAnnotation()
                                : tempNode.getAllDeniedAnnotation();
                        if (allAnnotation != null) {
                            permissions = needMethodMap.get(requestCode).getNeedAnnotation().getPermissions();
                        } else {
                            final PartAnnotation partAnnotation = isGranted
                                    ? tempNode.getOnGrantedAnnotation()
                                    : tempNode.getOnDeniedAnnotation();
                            permissions = partAnnotation.getPermissions();
                        }
                    } else {
                        final PartAnnotation partAnnotation = isGranted
                                ? tempNode.getOnGrantedAnnotation()
                                : tempNode.getOnDeniedAnnotation();
                        permissions = partAnnotation.getPermissions();
                    }
                    if (i == 0) {
                        subStart = new Label();
                    } else {
                        subStart = subNext;
                    }
                    if (i < infoSize - 1) {
                        subNext = new Label();
                    } else {
                        subNext = end;
                    }
                    methodVisitor.visitLabel(subStart);
                    if (i != 0) {
                        methodVisitor.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
                    }
                    methodVisitor.visitVarInsn(ALOAD, 2);
                    push(methodVisitor, permissions.size());
                    methodVisitor.visitTypeInsn(ANEWARRAY, "java/lang/String");
                    Iterator<String> iterator = permissions.iterator();
                    for (int j = 0; j < permissions.size() && iterator.hasNext(); j++) {
                        methodVisitor.visitInsn(DUP);
                        push(methodVisitor, j);
                        methodVisitor.visitLdcInsn(iterator.next());
                        methodVisitor.visitInsn(AASTORE);
                    }
                    methodVisitor.visitMethodInsn(INVOKESTATIC, STR_ARR_EQUAL_METHOD_OWNER,
                            STR_ARR_EQUAL_METHOD_NAME, STR_ARR_EQUAL_METHOD_DESC, false);
                    methodVisitor.visitJumpInsn(IFEQ, subNext);
                    Label label = new Label();
                    methodVisitor.visitLabel(label);
                    methodVisitor.visitVarInsn(ALOAD, 0);
                    methodVisitor.visitMethodInsn(getInvokeOpCode(tempNode.access), classNode.name, tempNode.name, tempNode.desc, false);
                    if (i != infoSize - 1 || index != codeSize - 1) {
                        methodVisitor.visitJumpInsn(GOTO, end);
                    }
                }
            }
            index++;
        }
        methodVisitor.visitLabel(end);
        methodVisitor.visitFrame(Opcodes.F_SAME, 0, null, 0, null);
        methodVisitor.visitInsn(RETURN);
        Label realEnd = new Label();
        methodVisitor.visitLabel(realEnd);
        methodVisitor.visitLocalVariable("this", getClassDesc(classNode.name), null, realStart, realEnd, 0);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_1, METHOD_PARAMETER_1_DESC, null, realStart, realEnd, 1);
        methodVisitor.visitLocalVariable(METHOD_PARAMETER_2, METHOD_PARAMETER_2_DESC, null, realStart, realEnd, 2);
        methodVisitor.visitMaxs(5, 3);
        methodVisitor.visitEnd();
    }

    private void push(final MethodVisitor methodVisitor, final int value) {
        if (value >= -1 && value <= 5) {
            methodVisitor.visitInsn(Opcodes.ICONST_0 + value);
        } else if (value >= Byte.MIN_VALUE && value <= Byte.MAX_VALUE) {
            methodVisitor.visitIntInsn(Opcodes.BIPUSH, value);
        } else if (value >= Short.MIN_VALUE && value <= Short.MAX_VALUE) {
            methodVisitor.visitIntInsn(Opcodes.SIPUSH, value);
        } else {
            methodVisitor.visitLdcInsn(value);
        }
    }

    private int getInvokeOpCode(int accessFlag) {
        return accessFlag == ACC_PRIVATE ? INVOKESPECIAL : INVOKEVIRTUAL;
    }

    private String getClassDesc(String clz) {
        return "L" + clz + ";";
    }


}
